// 导入所需的依赖包

mod logger;
mod sql_generator;

#[allow(unused_imports)]
use log::{debug, error, info, trace, warn};
// use futures::TryFutureExt;
use serde::Deserialize;
use sql_generator::TableInfo;
use sql_generator::{
    add_foreign_dependent, generate_random_sql, get_table_info_list,
    load_table_info_list_from_file, write_to_file,
};

use std::fs;
use std::io::Read;
use std::path::PathBuf;
use structopt::StructOpt;

use rand::Rng;
use tokio_postgres::Config as PgConfig;
use tokio_postgres::NoTls;

use std::path::Path;
use std::str::FromStr;
use tokio::time::Duration;
use tokio::time::Instant;
// use futures_util::future::try_future::TryFutureExt;
use futures_util::TryFutureExt;

mod config;
use config::CONFIG;

type LimitRange = (f64, u32);

#[derive(Debug, StructOpt)]
#[structopt(name = "sqlmaker", about = "A tool for processing SQL queries.")]
struct Opt {
    /// Configuration file
    #[structopt(short = "c", long = "config", parse(from_os_str))]
    config: PathBuf,
}

#[derive(Debug, Deserialize)]
struct ConfigFile {
    usercfg: UserCfg,
    appcfg: AppCfg,
}

#[derive(Debug, Deserialize)]
struct UserCfg {
    database_connection: String,
    sql_files_path: String,
    sql_output_path: String,
    sql_output_num: usize,
}

#[derive(Debug, Deserialize)]
struct AppCfg {
    sql_timeout_sec: u64,
    sql_exec_min_sec: f64,
    sql_output_count: usize,

    num_of_unique_values: u32, // 在数据库获取表信息时，唯一值unique_values的最大数量，默认20
    join_conditions_probability: f64, // 连接条件生成的概率 默认 0.7
    num_constraints_range: usize, // 定义约束条件数量的随机范围，默认最大 10
    constraint_probability: f64, // 有多大概率选择 IN 的表达式， 否则就是 is not NULL 默认 0.85
    null_probability: f64,     // 有多大概率生成 is not NULL， 否则就是 is NULL 默认 0.5
    num_values: usize,         // IN 表达式中，唯一值的数量，默认是 7
    select_fields_probability: f64, // 有多少的概率生成 COUNT(*)， 默认 0.4
    num_max_minormax: usize, // 不生成 SELECT COUNT(*) 的话，生成 SELECT [聚合函数]，生成[聚合函数]的最大数量 默认 10
    num_max_tables: usize,   // 随机选取一定数量的表 默认 12
    limit_ranges: String,    // 生成LIMIT子句, Vec<概率，最大值>

    table_info_filename: String,
    table_info_foreign_filename: String,
}

fn parse_limit_ranges(limit_ranges_str: &str) -> Vec<LimitRange> {
    limit_ranges_str
        .split(';')
        .map(|range| {
            let parts: Vec<&str> = range.split(',').collect();
            (
                parts[0].parse::<f64>().unwrap(),
                parts[1].parse::<u32>().unwrap(),
            )
        })
        .collect()
}
// fn test_extract_relationships_from_query() {
//     logger::init_logger();
//     let query = "SELECT MIN(chn.name) AS character_name, MIN(mi_idx.info) AS rating, MIN(t.title) AS complete_hero_movie FROM complete_cast AS cc, comp_cast_type AS cct1, comp_cast_type AS cct2, char_name AS chn, cast_info AS ci, info_type AS it2, keyword AS k, kind_type AS kt, movie_info_idx AS mi_idx, movie_keyword AS mk, name AS n, title AS t WHERE cct1.kind  = 'cast' AND cct2.kind  like '%complete%' AND chn.name  is not NULL and (chn.name like '%man%' or chn.name like '%Man%') AND it2.info  = 'rating' AND k.keyword  in ('superhero', 'marvel-comics', 'based-on-comic', 'tv-special', 'fight', 'violence', 'magnet', 'web', 'claw', 'laser') AND kt.kind  = 'movie' AND t.production_year  > 2000 AND kt.id = t.kind_id AND t.id = mk.movie_id AND t.id = ci.movie_id AND t.id = cc.movie_id AND t.id = mi_idx.movie_id AND mk.movie_id = ci.movie_id AND mk.movie_id = cc.movie_id AND mk.movie_id = mi_idx.movie_id AND ci.movie_id = cc.movie_id AND ci.movie_id = mi_idx.movie_id AND cc.movie_id = mi_idx.movie_id AND chn.id = ci.person_role_id AND n.id = ci.person_id AND k.id = mk.keyword_id AND cct1.id = cc.subject_id AND cct2.id = cc.status_id AND it2.id = mi_idx.info_type_id;
//     ";

//     let relationships = extract_relationships_from_query(&query);

//     for relationship in relationships {
//         println!("{:?}", relationship);
//     }
// }

fn update_config() {
    // 0. 解析命令行参数，读取配置文件...
    let opt = Opt::from_args();
    let config_path = opt.config;

    let mut config_contents = String::new();
    let mut file = fs::File::open(config_path).expect("Failed to open config file.");
    file.read_to_string(&mut config_contents)
        .expect("Failed to read config file.");

    let config_file: ConfigFile =
        toml::from_str(&config_contents).expect("Failed to parse config file.");

    // 概率参数获取 lazy_static
    // 更新配置 从配置文件中获取
    let mut config = CONFIG.write().unwrap();
    config.db_connection = config_file.usercfg.database_connection;
    config.sql_files_path = config_file.usercfg.sql_files_path;
    config.sql_output_path = config_file.usercfg.sql_output_path;
    config.sql_output_num = config_file.usercfg.sql_output_num;

    config.sql_timeout_sec = config_file.appcfg.sql_timeout_sec;
    config.sql_exec_min_sec = config_file.appcfg.sql_exec_min_sec;
    config.sql_output_count = config_file.appcfg.sql_output_count;

    config.num_of_unique_values = config_file.appcfg.num_of_unique_values; // 从配置文件中获取
    config.join_conditions_probability = config_file.appcfg.join_conditions_probability; // 从配置文件中获取
    config.num_constraints_range = config_file.appcfg.num_constraints_range; // 从配置文件中获取
    config.constraint_probability = config_file.appcfg.constraint_probability; // 从配置文件中获取
    config.null_probability = config_file.appcfg.null_probability; // 从配置文件中获取
    config.num_values = config_file.appcfg.num_values; // 从配置文件中获取
    config.select_fields_probability = config_file.appcfg.select_fields_probability; // 从配置文件中获取
    config.num_max_minormax = config_file.appcfg.num_max_minormax; // 从配置文件中获取
    config.num_max_tables = config_file.appcfg.num_max_tables; // 从配置文件中获取
    config.limit_ranges = parse_limit_ranges(&config_file.appcfg.limit_ranges);

    config.table_info_filename = config_file.appcfg.table_info_filename; // 从配置文件中获取
    config.table_info_foreign_filename = config_file.appcfg.table_info_foreign_filename;
    // 从配置文件中获取
}
// 主函数
fn main() {
    // 初始化日志记录器
    logger::init_logger();

    update_config();

    // 读取全局配置
    let config = CONFIG.read().unwrap();

    // 1. 获取表信息 table_info.json
    let table_info_list: Vec<TableInfo> =
        match load_table_info_list_from_file(&config.table_info_filename) {
            Ok(table_info_list) => {
                println!(
                    "检测到已有数据库信息文件，从 {} 文件中加载 table_info_list。",
                    &config.table_info_filename
                );
                table_info_list
            }
            Err(_) => {
                println!("首次运行，尚未收集数据库信息，正在查询数据库...");

                // 连接数据库并获取 table_info_list
                let table_info_list = tokio::runtime::Builder::new_multi_thread()
                    .worker_threads(2)
                    .enable_all()
                    .build()
                    .unwrap()
                    .block_on(get_table_info_list(&config.db_connection))
                    .expect("从数据库中获取表信息失败");

                // 将 table_info_list 序列化为 JSON 并写入文件
                let table_info_json = serde_json::to_string_pretty(&table_info_list).unwrap();
                write_to_file(&config.table_info_filename, &table_info_json).expect("写入文件失败");

                println!(
                    "已经集数据库信息，保存到 {} ...",
                    &config.table_info_filename
                );

                table_info_list
            }
        };

    debug!("table_info_list {:?}", table_info_list);

    // 2. 如果cfg::TABLE_INFO_FOREIGN_FILE_NAME不存在，为表信息添加外键依赖并从样例sql中学习
    if !Path::new(&config.table_info_foreign_filename).exists() {
        println!("正在为表信息添加外键依赖并从样例sql中学习...");
        add_foreign_dependent(&config.sql_files_path);
    }

    // 3. 从文件中加载添加了外键关系的 updated_tables.json
    let table_info_list_have_foreign_relation: Vec<TableInfo> =
        match load_table_info_list_from_file(&config.table_info_foreign_filename) {
            Ok(table_info_list) => {
                println!("从文件中加载 table_info_list_have_foreign_relation");
                table_info_list
            }
            Err(_) => {
                panic!("从文件中加载 table_info_list_have_foreign_relation 失败...");
            }
        };

    // 4. 生成随机 SQL
    let output_dir = Path::new(&config.sql_output_path);
    fs::create_dir_all(output_dir).expect("Failed to create output directory.");

    let rng = rand::thread_rng();
    let mut count = *&config.sql_output_count;
    let newline = if cfg!(windows) { "\r\n" } else { "\n" };

    while count < *&config.sql_output_num {
        let random_sql = generate_random_sql(&table_info_list_have_foreign_relation);

        // 在random_sql之前添加SET statement_timeout 语句，防止数据库超时
        let sql_with_timeout = format!(
            "SET statement_timeout TO {}; {}",
            *&config.sql_timeout_sec * 1000,
            random_sql
        );

        // 尝试执行 SQL
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .unwrap();
        let (status, duration) = runtime.block_on(async {
            let mut pg_config = PgConfig::from_str(&config.db_connection).unwrap();

            pg_config.connect_timeout(Duration::from_secs(*&config.sql_timeout_sec));
            let (client, connection) = match pg_config.connect(NoTls).await {
                Ok(result) => result,
                Err(_) => return ("FAILED", Duration::from_secs(0)),
            };

            let connection = connection.map_err(|e| eprintln!("connection error: {}", e));
            tokio::spawn(connection);

            let start = Instant::now();

            // let result = client.batch_execute(&random_sql).await;
            let result = tokio::time::timeout(
                Duration::from_secs(*&config.sql_timeout_sec), // 防止程序超时
                client.batch_execute(&sql_with_timeout),
            )
            .await;

            let end = Instant::now();
            let elapsed = end - start;

            match result {
                Ok(Ok(_)) => ("SUCCESS", elapsed),
                Ok(Err(_)) => ("FAILED", elapsed),
                Err(_) => ("TIMEOUT", elapsed),
            }
        });

        // 记录日志, 如果 SQL 执行成功，保存文件
        // 设置文件名
        // let file_name = if status == "SUCCESS" {
        //     let random_string: String = rng
        //         .clone()
        //         .sample_iter(&rand::distributions::Alphanumeric)
        //         .take(4)
        //         .map(char::from)
        //         .collect();
        //     format!("q{}_{}.sql", count + 1, random_string)
        // } else {
        //     "noname.sql".to_string()
        // };

        let file_name = if status == "SUCCESS" {
            let random_string: String = rng
                .clone()
                .sample_iter(&rand::distributions::Alphanumeric)
                .take(4)
                .map(char::from)
                .map(|c| c.to_ascii_lowercase()) // 将字符转换为小写
                .collect();
            format!("q{}_{}.sql", count + 1, random_string)
        } else {
            "noname.sql".to_string()
        };

        let min_sec:f64 = if config.sql_exec_min_sec > config.sql_timeout_sec as f64 {
            0.0
        } else {
            config.sql_exec_min_sec
        };

        // 如果 SQL 执行成功，保存文件
        if status == "SUCCESS" && duration.as_secs_f64() > min_sec {
            let file_path = output_dir.join(&file_name);
            fs::write(&file_path, random_sql.clone()).expect("Failed to write SQL to file.");
            count += 1;
        }

        // 记录日志
        info!(
            "[{}] [{}] [{}] [{}]",
            file_name,
            duration.as_secs_f64(),
            status,
            random_sql.replace(newline, " "),
        );
    }
}
